import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import math

fig = plt.figure()
ax = Axes3D(fig)


# 绘制2D点
def plt_point3d(point, c='r', marker='o'):
    x, y, z = point
    return ax.scatter(x, y, z, c=c, marker=marker)


def plt_point3d_wave(point, c='g', marker='o'):
    xw, yw, zw, ww = point
    if ww == 0:
        print("Error: w == 0")
    else:
        ax.scatter(xw / ww, yw / ww, zw / ww, c=c, marker=marker)


def plt_plane(line, color='b'):
    a, b, c, d = line
    # ax + by + cz + d = 0
    # (x,0,0)
    # (0,y,0)
    # (0,0,z)
    p1 = [1, 2, 3]
    p2 = [2, 1, 2]
    p3 = [3, 2, 1]
    if abs(a - 0) > 0.00001:
        p1[0] = -(b * 2 + c * 3 + d) / a
    if abs(b - 0) > 0.00001:
        p2[1] = -(a * 2 + c * 2 + d) / b
    if abs(c - 0) > 0.00001:
        p3[2] = -(a * 3 + c * 1 + d) / c

    ax.plot_trisurf(p1, p2, p3)


def plt_plane_theta(theta, d):
    # (cos(x)*cos(x), sin(x)*cos(x), sin(x), d)
    a = math.cos(theta) * math.cos(theta)
    b = math.cos(theta) * math.sin(theta)
    c = math.sin(theta)
    plt_plane((a, b, c, d))


if __name__ == '__main__':
    # (x, y, z)
    point3d = (1, 3, 1)
    plt_point3d(point3d)
    # (x, y, z, w)
    point3d_w = (2., 6., 4., 2.)
    plt_point3d_wave(point3d_w)
    # (0, 0, 0, 0)
    plane = (0, 0, 1.0, 10)
    # plt_plane(plane)

    plane = np.random.rand(4)
    plt_plane(plane)

    theta = math.pi / 2
    plt_plane_theta(theta, 10)

    plt.show()
